﻿# 导入必要的库
import gymnasium as gym  # 强化学习环境库
from stable_baselines3 import DQN  # 深度Q网络算法
from stable_baselines3.common.evaluation import evaluate_policy  # 用于评估模型性能

# 创建环境
# `render_mode="rgb_array"` 用于可视化训练过程（如果需要实时渲染，则需切换到非远程运行）
env = gym.make("LunarLander-v3", render_mode="human")

# 创建 DQN 模型
# - policy: "MlpPolicy" 表示使用多层感知机（MLP）作为策略网络
# - env: 强化学习的训练环境
# - learning_rate: 学习率，控制神经网络更新权重的步长
# - buffer_size: 经验回放缓冲区的大小，用于存储过去的经验以进行采样
# - learning_starts: 模型开始学习前需要收集的初始步骤数
# - batch_size: 每次训练时从经验缓冲区采样的样本数量
# - verbose: 日志的详细程度，0 表示不输出，1 表示输出进度，2 表示输出详细信息
model = DQN(
    policy="MlpPolicy",
    env=env,
    learning_rate=1e-3,
    buffer_size=100000,
    learning_starts=1000,
    batch_size=64,
    verbose=1,
)

# 开始训练
# `total_timesteps`: 总的训练步数
print("Starting training...")
model.learn(total_timesteps=50000)  # 训练 50000 个时间步

# 保存模型
model.save("dqn_lunarlander")

# 加载模型（如需要继续训练或评估）
# model = DQN.load("dqn_lunarlander", env=env)

# 评估模型性能
# 使用 `evaluate_policy` 评估训练好的策略，运行 10 个回合并返回平均奖励和标准差
mean_reward, std_reward = evaluate_policy(model, env, n_eval_episodes=10)
print(f"Mean reward: {mean_reward}, Std reward: {std_reward}")

# 关闭环境
env.close()
